#!/usr/bin/env python3
"""
Ingestion Agent 客户端脚本
支持对不同文件进行文档摄取处理
"""

import asyncio
import argparse
import os
import sys
from pathlib import Path
from typing import List, Optional
from dotenv import load_dotenv

load_dotenv()
# 添加项目根目录到Python路径
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from scorpio.core.common import get_logger, setup_logging, EmbeddingConfig, DatabaseConfig, PostgresConfigurationSettings
from scorpio.services.ingestion import IngestionService, ChunkConfig
from scorpio.agenthub.ingestion import IngestionAgent
from scorpio.core.database import VectorDatabaseManager
from scorpio.services.database import DocumentProcessor

# 初始化日志配置
LOG_LEVEL: str = os.getenv("LOG_LEVEL", "INFO")
LOG_DIR: str = os.getenv("LOG_DIR", "logs")
setup_logging(
    log_level=LOG_LEVEL,
    log_dir=LOG_DIR
)

logger = get_logger(__name__)

class IngestionAgentClient:
    """Ingestion Agent 客户端"""
    
    def __init__(self, db_manager: VectorDatabaseManager, embedding_config: EmbeddingConfig, chunk_config: ChunkConfig):
        """初始化客户端"""
        self.supported_formats = ['.pdf', '.docx', '.txt', '.md']
        self.embedding_config = embedding_config 
        logger.info(f"embedding_config: {self.embedding_config}")
        
        # 初始化数据库服务
        
        self.db_manager = db_manager
        self.document_processor = DocumentProcessor(db_manager, embedding_config)
        # 初始化Ingestion Service
        logger.info(f"chunk_config is: {chunk_config}")
        self.ingestion_service = IngestionService(chunk_config=chunk_config)
        self.shared =  {
            "context": {
                "step": "document",
                "db_manager": self.db_manager,
                "ingestion": self.ingestion_service,
                "processor": self.document_processor,
                "embedding_config": self.embedding_config,
                "ingestion": self.ingestion_service
            },
        }
    
    async def process(self, file_path: str, 
                    title: Optional[str] = None) -> dict:
        """处理单个文件"""
        
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        content_type = os.path.splitext(file_path)[1].lower()
        if content_type not in self.supported_formats:
            raise ValueError(f"不支持的文件格式: {content_type}")
        
        # 获取文件大小
        file_size = os.path.getsize(file_path)
        
        logger.info(f"开始处理文件: {file_path}")
        logger.info(f"文件大小: {file_size} bytes")
        self.shared["document"] = {
                "file_path": file_path,
                "title": title or os.path.basename(file_path),
            }
        try:
            agent=IngestionAgent(self.shared)
            # 使用Ingestion Service处理
            agent.create_flow()
            await agent.run()
            document = agent.shared["document"]
            return {
                'document_id': document["id"],
                'file_path': file_path,
                'title': title or os.path.basename(file_path),
                'status': 'success'
            }
            
        except Exception as e:
            logger.error(f"文件处理失败: {str(e)}")
            return {
                'document_id': document["id"],
                'file_path': file_path,
                'title': title or os.path.basename(file_path),
                'status': 'failed',
                'error': str(e)
            }
    


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Ingestion Agent 客户端")
    parser.add_argument("path", help="文件路径")
    parser.add_argument("--title", "-t", help="文档标题（可选）")
    parser.add_argument("--chunk-size", type=int, default=1000, help="分块大小")
    parser.add_argument("--chunk-overlap", type=int, default=200, help="分块重叠大小")
    parser.add_argument("--strategy", choices=["sentence", "paragraph", "fixed"], 
                        default="sentence", help="分块策略")
    parser.add_argument("--max-chunks", type=int, help="最大分块数量")
    
    args = parser.parse_args()
    
    # 创建分块配置
    chunk_config = ChunkConfig(
        chunk_size=args.chunk_size,
        chunk_overlap=args.chunk_overlap,
        strategy=args.strategy,
        max_chunks=args.max_chunks
    )
    print(os.getenv("PGHOST"), os.getenv("PGPORT"))
    db_config = DatabaseConfig(
        host=os.getenv("PGHOST"),
        port=os.getenv("PGPORT"),
        dbname=os.getenv("PGDATABASE"),
        user=os.getenv("PGUSER"),
        password=os.getenv("PGPASSWORD"),
        config_settings=PostgresConfigurationSettings(),
    )
    db_manager = VectorDatabaseManager(db_config)
    await db_manager.initialize()
    logger.info(f"Connect Vector database success with database config {db_config}...")
    # 初始化Agent客户端
    client = IngestionAgentClient(
                db_manager=db_manager,
                embedding_config=EmbeddingConfig(), 
                chunk_config=chunk_config)
    try:
        if os.path.isfile(args.path):
            result = await client.process(
                args.path, 
                args.title 
            )
            print(f"\n##\n处理结果: {result}")
        else:
            print(f"路径不存在: {args.path}")
            sys.exit(1)
    except Exception as e:
        logger.error(f"客户端执行失败: {str(e)}")
        sys.exit(1)
    
if __name__ == "__main__":
    asyncio.run(main())